module ProxyPump
  class ThreadPool

    class ForceShutdownError < RuntimeError ; end

    attr_reader :min, :max

    def initialize(min, max = nil)
      @min       = min
      @max       = max || min

      @lock      = Mutex.new

      @tasks     = Queue.new
      @workers   = Array.new

      @spawned   = 0
      @idle      = 0
      @shutdown  = false
      sync { min.times { _spawn_thread } }
    end

    # Show how many tasks in queue.
    #
    def backlog
      @tasks.size
    end

    def spawned
      sync { @spawned }
    end
    alias size spawned

    def idle
      sync { @idle }
    end

    # Add task to pool.
    #
    def schedule(*args, &block)
      tap do
        sync do
          raise "Unable to add more tasks while shutting down" if @shutdown
          _schedule(*args, &block)
          # Try to spawn a new thread to run task if not enough.
          _spawn_thread if _short_handed?
        end
      end
    end
    alias << schedule

    def shutdown(timeout = -1)
      workers = sync do
        @shutdown = true
        @workers.dup
      end
      timeout = timeout.to_i
      if timeout <= 0
        # Gracefully shutdown, wait for all threads to finish.
        workers.size.times { _schedule { throw :exit } }
        workers.map(&:join)
      else
        # Wait for threads to finish after specified timeout seconds.
        # If there are unfinished threads, then force kill them.
        timeout.times do
          workers.reject! do |t|
            # If the time limit expires, nil will be returned,
            # otherwise thread is returned, then delete it.
            #
            t.join 1
          end
          if workers.empty?
            break
          else
            sleep 1
          end
        end
        # If pool is no empty, then force shutdown by raising exception.
        workers.each { |t| t.raise ForceShutdownError }

        workers.each { |t| t.join 5 }
      end
    end

    # If there are too many idle threads, tell one to exit.
    # If +force+ is true, then a trim request is requested event if all threads are busy.
    #
    def trim(force = false)
      tap do
        sync do
          if (force || @idle > 0) && @spawned > @min
            th = @workers.detect { |t| t[:idle] }
            if th.exit
              @workers.delete th
              @spawned -= 1
            end
          end
        end
      end
    end
    alias layoff trim

    # IF there are dead threads in pool, clean up them, keep the pool healthy.
    #
    def reap
      sync do
        dead_workers = @workers.reject(&:alive?)
        dead_workers.each do |t|
          t.kill # Same as #exit or #terminate ?
          @spawned -= 1
        end
        @workers.reject! { |t| dead_workers.include? t }
      end
    end
    alias clean_up reap

    protected

      def sync
        @lock.synchronize { yield }
      end

    private

      # Create a new thread to receive task.
      #
      def _spawn_thread
        @spawned += 1
        thread = Thread.new do
          Thread.current.name = 'Thread #%03i' % @spawned
          catch(:exit) do
            loop do
              @idle += 1
              Thread.current[:idle] = true
              task, args = @tasks.pop
              @idle -= 1
              Thread.current[:idle] = false
              # Ignore exceptions to prevent the thread to be killed.
              begin
                task.call *args
              rescue => e
                STDERR.puts "#{Thread.current.name} error: #{e.message} (#{e.class})"
                STDERR.puts e.backtrace
                e
              end
            end # loop
          end # catch

          # If go here, indicate that the pool is shutting down, let's remove the thread.
          sync do
            @spawned -= 1
            @workers.delete thread
          end

        end # Thread.new

        @workers << thread
        thread
      end

      def _schedule(*args, &block)
        @tasks << [block, args]
      end

      def _short_handed?
        @idle < @tasks.size && @spawned < @max
      end

  end # ThreadPool
end # ProxyPump
